package main

import (
	"crypto/rand"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"time"

	"github.com/spf13/cobra"
	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/x509"
)

func GenerateCertKeySM2(host, commonName string, alternateIPs []net.IP, alternateDNS []string, certFileName, privKeyFileName string) error {
	priv, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return (err)
	}

	max := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, _ := rand.Int(rand.Reader, max)
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Issuer:       pkix.Name{},
		Subject: pkix.Name{
			CommonName:   commonName,
			Organization: []string{"Hynet"},
			Country:      []string{"CN"},
			ExtraNames: []pkix.AttributeTypeAndValue{
				{
					Type:  []int{2, 5, 4, 42},
					Value: "Gopher",
				},
				{
					Type:  []int{2, 5, 4, 6},
					Value: "NL",
				},
			},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour * 24 * 365),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
		SignatureAlgorithm:    x509.SM2WithSM3,
	}

	if ip := net.ParseIP(host); ip != nil {
		template.IPAddresses = append(template.IPAddresses, ip)
	} else {
		template.DNSNames = append(template.DNSNames, host)
	}

	template.IPAddresses = append(template.IPAddresses, alternateIPs...)
	template.DNSNames = append(template.DNSNames, alternateDNS...)

	publicKey := priv.Public().(*sm2.PublicKey)
	certificateToPem, err := x509.CreateCertificateToPem(&template, &template, publicKey, priv)
	if err != nil {
		return (err)
	}

	file, err := os.Create(certFileName)
	if err != nil {
		return (err)
	}
	defer file.Close()
	file.Write(certificateToPem)

	privateKey, err := x509.MarshalSm2PrivateKey(priv, []byte("12345678"))
	if err != nil {
		return (err)
	}
	block := pem.Block{
		Type:    "SM2 PRIVATE KEY",
		Headers: nil,
		Bytes:   privateKey,
	}
	file, err = os.Create(privKeyFileName)
	if err != nil {
		return (err)
	}
	defer file.Close()
	err = pem.Encode(file, &block)
	if err != nil {
		return (err)
	}
	return nil
}

var genKeyCmd = &cobra.Command{
	Use:   "new",
	Short: "公私钥产生命令",
	Long: `例如：gmTool.exe new --private [私钥文件] --public [公钥证书文件名] --name 汇元网密钥
	`,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("generate key called")

		priFileName, err := cmd.Flags().GetString("private")
		if err != nil {
			fmt.Printf("Error: %v\n", err.Error())
			return
		}
		if len(priFileName) == 0 {
			fmt.Printf("Error: 私钥文件名不能为空\n")
			return
		}
		if FileExist(priFileName) {
			fmt.Printf("Error: 私钥文件 %v 已经存在，请妥善备份好原有的私钥再创建新的！\n", priFileName)
			return
		}

		pubFileName, err := cmd.Flags().GetString("public")
		if err != nil {
			fmt.Printf("Error: %v\n", err.Error())
			return
		}
		if len(pubFileName) == 0 {
			fmt.Printf("Error: 证书文件名不能为空\n")
			return
		}
		if FileExist(pubFileName) {
			fmt.Printf("Error: 公钥文件 %v 已经存在，请妥善备份好原有的私钥再创建新的！\n", pubFileName)
			return
		}

		commonName, err := cmd.Flags().GetString("name")
		if err != nil {
			fmt.Printf("Error: %v\n", err.Error())
			return
		}
		if len(commonName) == 0 {
			fmt.Printf("Error: 证书名称不能为空\n")
			return
		}

		ip := []byte("127.0.0.1")
		alternateDNS := []string{"localhost"}
		err = GenerateCertKeySM2("localhost", commonName, []net.IP{net.ParseIP(string(ip))}, alternateDNS, pubFileName, priFileName)
		if err != nil {
			fmt.Printf("Error: %v\n", err.Error())
			return
		}
		fmt.Printf("生成成功，请在本地妥善保存好生成的私钥文件：%v, （公钥证书：%v）\n", priFileName, pubFileName)
	},
}

func init() {
	rootCmd.AddCommand(genKeyCmd)

	// Here you will define your flags and configuration settings.
	genKeyCmd.Flags().StringP("private", "p", "sm2ca.key", "私钥PEM文件名称")
	genKeyCmd.Flags().StringP("public", "g", "sm2ca.cer", "公钥证书文件名称")
	genKeyCmd.Flags().StringP("name", "n", "", "证书名称")
}
